#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
@author: liang kang
@contact: gangkanli1219@gmail.com
@time: 1/2/18 11:28 AM
@desc: tools to generate crowd counting data
"""
import os

import cv2
import numpy as np

from utils.basic import get_file_name


def generate_gaussian_kernel(shape, sigma):
    """
    generate 2d gaussian kernel, default mean is zero
    Parameters
    ----------
    shape
    sigma

    Returns
    -------

    """
    rows, cols = shape

    def get_gaussian_fct(size):
        kernel1d = np.linspace(0, size, size)
        kernel1d -= (size / 2)
        kernel1d = kernel1d ** 2
        kernel1d /= (2 * sigma ** 2)
        return np.exp(-kernel1d)

    kernel = np.outer(get_gaussian_fct(rows), get_gaussian_fct(cols))
    return kernel / np.sum(kernel)


def check_all_crowd_counting(file_path, lists, params):
    """
    a filter for traverse file system

    Parameters
    ----------
    file_path: file absolution path
    lists: a series parameters needed, it likes:
                   0: file list of all the required file,
                   1: dict of the class name as the key and the number of
                   the class as the value
    params: some additional parameters, it likes:
                    0: class name list, the required class names
                    1: a string for split the different files

    Returns
    -------
    lists
    """
    if not file_path.endswith('.jpg'):
        return lists
    root = os.path.dirname(file_path)
    name = get_file_name(file_path)
    if os.path.exists(os.path.join(root, name + '.txt')):
        lists[0].append(os.path.join(root, name + '.txt') + params + file_path)
    return lists


def check_dat_and_img_file(file_path, lists, params):
    """
    a filter for traverse file system

    Parameters
    ----------
    file_path: file absolution path
    lists: a series parameters needed, it likes:
                   0: file list of all the required file,
                   1: dict of the class name as the key and the number of
                   the class as the value
    params: some additional parameters, it likes:
                    0: class name list, the required class names
                    1: a string for split the different files

    Returns
    -------
    lists
    """
    if file_path.split('.')[-1] in ['jpg', 'JPG', 'png', 'PNG']:
        shape = cv2.imread(file_path).shape
        if not (shape[0] == shape[1] == 1280):
            return lists
        if os.path.exists(file_path + '.dat'):
            lists[0].append(file_path + '.dat' + params + file_path)
    return lists


def decode_dat_file(file_name):
    """
    解析dat文件，并返回人的坐标点

    Parameters
    ----------
    file_name: file absolution path

    Returns
    -------

    """
    with open(file_name, 'rt') as dat_file:
        data = {}
        data_list = list(filter(lambda x: x != '\n', dat_file.readlines()))
        data_list = list(map(lambda x: x.strip(), data_list))
        data_list = data_list[2:]
        index = 0
        for data_buf in data_list:
            if data_buf == '0' or data_buf == '1':
                index += 1
                data[index] = []
            else:
                data[index].append(list(map(lambda x: int(x), data_buf.split('\t'))))
        points = []
        for _, value in data.items():
            if len(value) == 0:
                continue
            point_array = np.array(value)
            point = np.mean(point_array, axis=0)
            points.append(point.tolist())
        return points


def get_density_map(shape, points, kernel_size=15, sigma=4.0):
    """
    create a density map for crowd counting

    Parameters
    ----------
    shape: shape of density map, must be 2-D
    points: points of human
    kernel_size: base size of the gaussian filter
    sigma:

    Returns
    -------

    """
    density_map = np.zeros(shape[:2])
    if len(points) == 0:
        return density_map
    half_ks = kernel_size // 2
    for point in points:
        h = min(shape[0] - 1, max(0, point[0]))
        w = min(shape[1] - 1, max(0, point[1]))
        ks_h = min(half_ks, h, shape[0] - h - 1) * 2 + 1
        ks_w = min(half_ks, w, shape[1] - w - 1) * 2 + 1
        kernel = generate_gaussian_kernel((ks_h, ks_w), sigma)
        seat_h = h - ks_h // 2
        seat_w = w - ks_w // 2
        density_map[seat_h:(seat_h + ks_h), seat_w:(seat_w + ks_w)] += kernel
    return density_map
